from grid_env_ideal_obs_repeat_task import *
from grid_agent import *
from checkpoint_utils import *
from maze_factory import *
from replay_config import *
import argparse
import json
import sys
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.patches import Circle
from sklearn.manifold import TSNE
import random
from sklearn.decomposition import PCA
from matplotlib.animation import FuncAnimation
from sklearn.cluster import KMeans
import threading
from matplotlib.ticker import FuncFormatter

# analysis of phase space

def progress_bar(current, total, barLength = 100):
    percent = float(current) * 100 / total
    arrow = '-' * int(percent/100 * barLength - 1) + '>'
    spaces = ' ' * (barLength - len(arrow))

    print('Progress: [%s%s] %d %%' % (arrow, spaces, percent), end='\r')
    sys.stdout.flush()

@jit
def get_action(y):
    return jnp.argmax(y)
get_action_vmap = jax.vmap(get_action)

def ivf():

    """ parse arguments
    """
    rpl_config = ReplayConfig()

    parser = argparse.ArgumentParser()
    parser.add_argument("--model_pth", type=str, default=rpl_config.model_pth)
    parser.add_argument("--map_size", type=int, default=rpl_config.map_size)
    parser.add_argument("--task_pth", type=str, default=rpl_config.task_pth)
    parser.add_argument("--log_pth", type=str, default=rpl_config.log_pth)
    parser.add_argument("--nn_size", type=int, default=rpl_config.nn_size)
    parser.add_argument("--nn_type", type=str, default=rpl_config.nn_type)
    parser.add_argument("--show_kf", type=str, default=rpl_config.show_kf)
    parser.add_argument("--visualization", type=str, default=rpl_config.visualization)
    parser.add_argument("--video_output", type=str, default=rpl_config.video_output)
    parser.add_argument("--life_duration", type=int, default=rpl_config.life_duration)
    parser.add_argument("--start_i", type=int, default=rpl_config.start_i)
    parser.add_argument("--end_i", type=int, default=rpl_config.end_i)

    args = parser.parse_args()

    rpl_config.model_pth = args.model_pth
    rpl_config.map_size = args.map_size
    rpl_config.task_pth = args.task_pth
    rpl_config.log_pth = args.log_pth
    rpl_config.nn_size = args.nn_size
    rpl_config.nn_type = args.nn_type
    rpl_config.show_kf = args.show_kf
    rpl_config.visualization = args.visualization
    rpl_config.video_output = args.video_output
    rpl_config.life_duration = args.life_duration
    rpl_config.start_i = args.start_i
    rpl_config.end_i = args.end_i

    """ load model
    """
    params = load_weights(rpl_config.model_pth)
    # 定义一个函数，用于生成随机权重
    def init_weights_r(key, shape):
        return jax.random.normal(key, shape)
    # 生成一个随机的 PRNGKey
    key = jax.random.PRNGKey(np.random.randint(0, 1000))
    # key = jax.random.PRNGKey(4512)
    # 使用 tree_map 遍历 params 对象，并使用 init_weights 函数生成随机权重
    random_params = jax.tree_map(lambda x: init_weights_r(key, x.shape), params)
    # params = random_params

    # get elements of params
    tree_leaves = jax.tree_util.tree_leaves(params)
    for i in range(len(tree_leaves)):
        print("shape of leaf ", i, ": ", tree_leaves[i].shape)

    bias1 = np.array(tree_leaves[0])
    mat1 = np.array(tree_leaves[1])
    print("mat1.shape: ", mat1.shape)
    print("bias1: ", bias1)
    bias2 = np.array(tree_leaves[2])
    mat2 = np.array(tree_leaves[3])

    mat_obs = np.array(tree_leaves[1])[rpl_config.nn_size:rpl_config.nn_size+9,:]
    mat_intr = np.array(tree_leaves[1])[0:rpl_config.nn_size,:]
    print("mat_obs.shape: ", mat_obs.shape)

    """ create agent
    """
    if rpl_config.nn_type == "vanilla":
        model = RNN(hidden_dims = rpl_config.nn_size)
    elif rpl_config.nn_type == "gru":
        model = GRU(hidden_dims = rpl_config.nn_size)

    # check if param fits the agent
    if rpl_config.nn_type == "vanilla":
        assert params["params"]["Dense_0"]["kernel"].shape[0] == rpl_config.nn_size + 10

    n_samples = 10000
    k1 = npr.randint(0, 1000000)
    rnn_state = model.initial_state_rnd(n_samples, k1)
    rnn_state_old = rnn_state.copy()
    diff = jnp.abs(rnn_state - rnn_state_old)
    rnn_state_old = rnn_state.copy()
    diff_norm = jnp.linalg.norm(diff, axis=1)
    diff_norm_old = diff_norm.copy()
    norm_std = diff_norm.copy()

    rnn_state_trajectory = []
    intr_field = []

    rnn_action = np.zeros((n_samples, 4))
    # rnn_actions = []

    for t in range(rpl_config.life_duration):

        progress_bar(t, rpl_config.life_duration)

        intr_vector = np.dot(rnn_state, mat_intr) + bias1
        intr_vector = np.tanh(intr_vector)
        intr_field.append(intr_vector - rnn_state)

        rnn_state_trajectory.append(np.array(rnn_state).copy())

        # rnn_action = np.dot(rnn_state, mat2) + bias2
        # rnn_actions.append(rnn_action)
        
        """ model forward 
        """
        rnn_state = intr_vector
        
    print(rnn_state.shape)
    print(norm_std.shape)
    rnn_state_np = np.array(rnn_state)

    # 将 rnn_state_trajectory 展开成 rnn_state_np 的形状
    rnn_state_trajectory_np = np.array(rnn_state_trajectory)
    rnn_state_trajectory_np = rnn_state_trajectory_np.reshape(-1, rnn_state_trajectory_np.shape[-1])
    print("shape of rnn_state_trajectory_np: ", rnn_state_trajectory_np.shape)

    # compute actions
    rnn_actions0 = np.dot(rnn_state_trajectory_np, mat2) + bias2
    print("shape of rnn_actions: ", rnn_actions0.shape)
    actions0 = get_action_vmap(jnp.array(rnn_actions0))
    actions0 = np.array(actions0)
    print("shape of actions: ", actions0.shape)

    # 将 intr_field 展开成 rnn_state_np 的形状
    intr_field_np = np.array(intr_field)
    intr_field_np = intr_field_np.reshape(-1, intr_field_np.shape[-1])
    print("shape of intr_field_np: ", intr_field_np.shape)

    # 对 rnn_state_np 进行 PCA
    intrinsic_pca = PCA()
    intrinsic_pca.fit(rnn_state_trajectory_np)

    # # 打印 variance ratio
    # print(intrinsic_pca.explained_variance_ratio_)
    # explained_variance_ratio_sum = np.zeros_like(intrinsic_pca.explained_variance_ratio_)
    # for i in range(len(intrinsic_pca.explained_variance_ratio_)):
    #     explained_variance_ratio_sum[i] = np.sum(intrinsic_pca.explained_variance_ratio_[0:i+1])
    # plt.bar(range(128), explained_variance_ratio_sum)
    # plt.show()

    rnn_state_trajectory_np_pca = intrinsic_pca.transform(rnn_state_trajectory_np)

    print("shape of rnn_state_trajectory_np_pca: ", rnn_state_trajectory_np_pca.shape)

    intr_pca = intrinsic_pca.transform(intr_field_np + intrinsic_pca.mean_)

    compute_dim0 = 0
    compute_dim1 = 1
    compute_dim2 = 2

    view_dim0 = compute_dim0
    view_dim1 = compute_dim1
    view_dim2 = compute_dim2

    compute_dims = [compute_dim0, compute_dim1, compute_dim2]
    # compute_dims = [i for i in range(60)]

    # 将除了指定维度以外的所有维度置零
    mask = np.zeros(rnn_state_trajectory_np_pca.shape[1], dtype=bool)
    mask[compute_dims] = True
    rnn_state_trajectory_np_pca_masked = rnn_state_trajectory_np_pca.copy()
    rnn_state_trajectory_np_pca_masked[:, ~mask] = 0

    # # 将 pca 空间的点，逆变换回原始相空间
    HS_trajectory_1 = intrinsic_pca.inverse_transform(rnn_state_trajectory_np_pca_masked)

    # compute actions
    rnn_actions1 = np.dot(HS_trajectory_1, mat2) + bias2
    print("shape of rnn_actions: ", rnn_actions1.shape)
    actions1 = get_action_vmap(jnp.array(rnn_actions1))
    actions1 = np.array(actions1)
    print("shape of actions: ", actions1.shape)

    # 计算主成分空间中的内禀向量场
    intr_field1 = []
    for i in range(HS_trajectory_1.shape[0]):
        intr_vector = np.dot(HS_trajectory_1[i], mat_intr) + bias1
        intr_vector = np.tanh(intr_vector)
        intr_field1.append(intr_vector - HS_trajectory_1[i])
    # 将内禀向量场投影到主成分空间
    intr_pca_1 = intrinsic_pca.transform(np.array(intr_field1) + intrinsic_pca.mean_)

    # 比较 intr_pca_1 和 intr_pca 的差异
    print("intr_pca_1 - intr_pca: ", np.linalg.norm(intr_pca_1[:,0:3] - intr_pca[:,0:3]))

    # 创建KMeans对象，指定聚类数为4
    kmeans = KMeans(n_clusters=4)
    # 对rnn_state_np_pca进行聚类
    kmeans.fit(rnn_state_np)
    # 获取聚类中心的坐标
    cluster_centers = kmeans.cluster_centers_
    cluster_centers_pca = intrinsic_pca.transform(cluster_centers)

    # fig = plt.figure()
    # ax = fig.add_subplot(111, projection='3d')
    
    # rnn_action_diff = np.linalg.norm(rnn_actions0 - rnn_actions1, axis=1)
    # action_diff = np.abs(actions0 - actions1)
    # print("shape of rnn_action_diff: ", rnn_action_diff.shape)
    # action_diff_bin = np.where(action_diff > 0, True, False)

    # rnn_action_diff = rnn_action_diff/np.max(rnn_action_diff)

    # total_ = 0
    # diff_ = 0

    # index_list = []
    # color_list = []

    # for i in range(rnn_state_trajectory_np_pca.shape[0] - 1):
    #     if i % 1000 == 0:
    #         index_list.append(i)
    #         total_ += 1
    #         if action_diff_bin[i]:
    #             color = 'red'
    #             diff_ += 1
    #         else:
    #             color = 'blue'
    #         color_list.append(color)

    # ax.scatter(rnn_state_trajectory_np_pca[index_list,view_dim0], rnn_state_trajectory_np_pca[index_list,view_dim1], rnn_state_trajectory_np_pca[index_list,view_dim2], s=rnn_action_diff[index_list]*100, color=color_list)
    # plt.show()

    # print("diff ratio : ", diff_/total_)

    def get_action_representative_ratio(k):

        compute_dims = [i for i in range(k)]

        # 将除了指定维度以外的所有维度置零
        mask = np.zeros(rnn_state_trajectory_np_pca.shape[1], dtype=bool)
        mask[compute_dims] = True
        rnn_state_trajectory_np_pca_masked = rnn_state_trajectory_np_pca.copy()
        rnn_state_trajectory_np_pca_masked[:, ~mask] = 0

        # # 将 pca 空间的点，逆变换回原始相空间
        HS_trajectory_1 = intrinsic_pca.inverse_transform(rnn_state_trajectory_np_pca_masked)

        # compute actions
        rnn_actions1 = np.dot(HS_trajectory_1, mat2) + bias2
        actions1 = get_action_vmap(jnp.array(rnn_actions1))
        actions1 = np.array(actions1)

        rnn_action_diff = np.linalg.norm(rnn_actions0 - rnn_actions1, axis=1)
        action_diff = np.abs(actions0 - actions1)
        action_diff_bin = np.where(action_diff > 0, True, False)

        n_diff = np.sum(action_diff_bin)
        total_diff = np.sum(rnn_action_diff)

        print("n_diff, action_diff_bin.shape[0]: ", n_diff, action_diff_bin.shape[0])

        return n_diff/action_diff_bin.shape[0], total_diff

    # print("n_diff, total_diff: ", get_action_representative_ratio(20))

    n_diff_list = []
    total_diff_list = []
    for i in range(3,129):
        progress_bar(i, 129-3)
        n_diff, total_diff = get_action_representative_ratio(i)
        n_diff_list.append(n_diff)
        total_diff_list.append(total_diff)
    # 将 n_diff_list 和 total_diff_list 绘制成曲线，分别画到两个独立的子图中
    fig, ax = plt.subplots(2, 1)
    ax[0].plot(np.array(n_diff_list))
    ax[1].plot(np.array(total_diff_list))

    # 设置横坐标的批注
    ax[0].set_xticks(np.arange(len(n_diff_list)))
    ax[0].set_xticklabels(np.arange(3, len(n_diff_list)+3))
    ax[1].set_xticks(np.arange(len(total_diff_list)))
    ax[1].set_xticklabels(np.arange(3, len(total_diff_list)+3))

    # 将纵坐标转换成百分比
    def to_percent(y, position):
        return str(round(y*100, 2)) + '%'
    ax[0].yaxis.set_major_formatter(FuncFormatter(to_percent))


    # 给两个子图分别命名
    ax[0].set_title("action_diff")
    ax[1].set_title("action_vector_diff")
    plt.show()


if __name__ == "__main__":
    
    ivf()